# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
import os
from concurrent.futures import ProcessPoolExecutor
import functools

from absl import (app, flags)
import numpy
import pandas

from coin.base.datetime_util import (convert_string_to_dates, iterate_date)
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.mm.subscription import get_subscription_key
from coin.proto.coin_order_gateway_pb2 import OrderEvent
from coin.exchange.base.order_gateway_logger import parse_file_name
from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.support.proto_log.logic.util import (run_from_og_log_archive,
                                               enumerate_og_log_into_oginfos)
from xunkemgmt.support.accounting.logic.query_request import get_all_requests
from coin.strategy.accounting.fee import get_pnl_fee
from coin.strategy.mm.tool.log_pta_dumper import (
    PtaStrategy,)
from coin.support.feed_tool.feed_stats.logic.util import (query_quote_to_usd_exchange_rate,
                                                          gen_engine)

FLAGS = flags.FLAGS


def query_symbol_vwap(trading_date, symbol):
  sql_tmpl = """
SELECT
exchange,
machine,
trading_date,
symbol,
vwap
FROM coin_feed_stat_20180514.StatsView
WHERE trading_date = {trading_date}
AND machine = 'feed-01.ap-northeast-1.aws'
AND exchange = 'Binance'
AND symbol = '{symbol}' 
;
"""

  trading_date_str = trading_date.strftime('%Y%m%d')
  sql = sql_tmpl.format(trading_date=trading_date_str, symbol=symbol)

  engine = gen_engine()
  feed_stats = pandas.read_sql(sql, engine)

  if len(feed_stats) >= 1:
    vwap = feed_stats.iloc[0]['vwap']
  else:
    # Try the next day.
    trading_date = trading_date + datetime.timedelta(days=1)
    vwap = query_symbol_vwap(trading_date, symbol)
  return vwap


class FeedDumper(object):
  def __init__(self, products):
    self._products = products
    self._bbos = {product: [] for product in products}

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      subscription_key = get_subscription_key(product)
      if book_builder_name == subscription_key:
        book_builder.subscribe(product, functools.partial(self.on_book, product))

  def on_book(self, product, book):
    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    self._bbos[product].append({
        'timestamp': book.timestamp, 'bid0p': bid0_price, 'ask0p': ask0_price
    })

  def get_product_df(self, product):
    if product not in self._products:
      return None
    if len(self._bbos[product]) == 0:
      return None
    df = pandas.DataFrame(self._bbos[product])
    df['datetime'] = pandas.to_datetime(df['timestamp'], unit='ns')
    df = df.set_index('datetime')

    return df


class OrderEventDumper(PtaStrategy):
  def __init__(self, market_type, exchange, trading_date, machine, owner):
    super().__init__(market_type, exchange, trading_date, machine)
    self._owner = owner

  def generate_fill_csv(self, root_path):
    quotes = set()
    for symbol in self._rows.keys():
      quote = symbol.split('-')[1]
      quotes.add(quote)

    quote_df_map = {}
    coin_quote_list = ['BTC', 'ETH']
    for quote in quotes:
      if quote not in coin_quote_list:
        continue
      usdt_product = BinanceProduct.FromStr('%s-USDT' % quote)
      hours = 24.5
      products = [usdt_product]

      strategy = FeedDumper(products=products)
      start_time = self._trading_date

      run_from_archive(products,
                       strategy.on_book_reset,
                       start_time,
                       start_time + datetime.timedelta(hours=hours),
                       machine=['feed-01.ap-northeast-1.aws'],
                       worker_id=['1'],
                       feed_checker_config=None,
                       skip_failed_product=False)

      df = strategy.get_product_df(usdt_product)
      quote_df_map[quote] = df

    for symbol in self._rows.keys():
      self._stat_df[symbol] = pandas.DataFrame(self._rows[symbol], columns=self._header)

      fill_df = self._process_oe_df(self._stat_df[symbol], symbol, quote_df_map)
      output_csv_path = self._get_output_csv_path(symbol, root_path)
      fill_df.to_csv(output_csv_path, float_format='%12.6f', index=False)

  def _process_oe_df(self, oe_df, symbol, quote_df_map):
    assert len(oe_df['symbol'].unique()) == 1
    fill_df = oe_df.loc[oe_df['event_type'] == OrderEvent.ORDER_FILLED].reset_index(drop=True)

    if len(fill_df) == 0:
      print('Fill empty: ',
            self._trading_date,
            self._machine,
            self._market_type,
            self._exchange,
            self._owner)
      return fill_df

    agg_idx = ~numpy.isclose(fill_df['fill_price'], fill_df['price'], atol=0)
    fill_df['agg'] = agg_idx
    fill_df['datetime'] = pandas.to_datetime(fill_df['timestamp'], unit='ns')
    fill_df['datetime_index'] = pandas.to_datetime(fill_df['timestamp'], unit='ns')
    fill_df = fill_df.set_index('datetime_index')
    fill_df['exchange'] = self._exchange
    fill_df['account'] = self._owner
    fill_df['transaction_date'] = self._trading_date.strftime('%Y-%m-%d')
    # fill_df['base'] = fill_df['symbol'].str.split('-').str[0]
    # fill_df['quote'] = fill_df['symbol'].str.split('-').str[1]
    fill_df['base'] = symbol.split('-')[0]
    fill_df['quote'] = symbol.split('-')[1]
    fill_df['side'] = fill_df.apply(self._translate_side, axis=1)
    fill_df['value'] = fill_df['fill_price'] * fill_df['fill_qty']

    fill_df['fee'] = fill_df.apply(self._calc_fee, axis=1)
    fill_df['fee_coin'] = fill_df['quote']

    quote = symbol.split('-')[1]
    if quote in quote_df_map:
      feed_df = quote_df_map[quote]
      nearest_df = feed_df.reindex(fill_df.index, method='nearest')
      nearest_df['coin_usd_price'] = (nearest_df['bid0p'] + nearest_df['ask0p']) / 2
      fill_df = pandas.merge(fill_df, nearest_df, how='inner', left_index=True, right_index=True)
      fill_df['usd_value'] = fill_df['value'] * fill_df['coin_usd_price']
      fill_df['usd_fee'] = fill_df['fee'] * fill_df['coin_usd_price']
    else:
      to_usd_rate = self.get_quote_to_usd_rate(symbol)

      fill_df['coin_usd_price'] = to_usd_rate
      fill_df['usd_value'] = fill_df['value'] * to_usd_rate
      fill_df['usd_fee'] = fill_df['fee'] * to_usd_rate

    out_fill_df = fill_df[[
        'datetime',
        'exchange',
        'account',
        'transaction_date',
        'base',
        'quote',
        'side',
        'fill_price',
        'fill_qty',
        'value',
        'coin_usd_price',
        'usd_value',
        'fee',
        'usd_fee',
        'fee_coin'
    ]]
    out_fill_df.reset_index(drop=True)

    return out_fill_df

  def get_quote_to_usd_rate(self, symbol):
    fx_quote_list = ['KRW', 'EUR', 'JPY']
    coin_quote_list = ['BTC', 'ETH']
    quote = symbol.split('-')[1]
    if quote in ['USD', 'USDT']:
      return 1
    elif quote in fx_quote_list:
      fx_rate = query_quote_to_usd_exchange_rate(fx_quote_list)
      return fx_rate[quote]
    elif quote in coin_quote_list:
      query_symbol = '%s-USDT' % quote
      vwap = query_symbol_vwap(self._trading_date, query_symbol)
      return vwap
    else:
      raise ValueError(symbol, self._trading_date, self._machine, self._exchange, self._owner)

  def _translate_side(self, fill_df):
    if fill_df['side'] == 1:
      return 'buy'
    elif fill_df['side'] == 2:
      return 'sell'
    else:
      raise ValueError(fill_df['side'],
                       self._trading_date,
                       self._machine,
                       self._exchange,
                       self._owner)

  def _calc_fee(self, fill_df):
    if fill_df['agg']:
      return get_pnl_fee(maker_pq=0,
                         taker_pq=fill_df['value'],
                         market_type=self._market_type,
                         exchange=self._exchange)
    else:
      return get_pnl_fee(maker_pq=fill_df['value'],
                         taker_pq=0,
                         market_type=self._market_type,
                         exchange=self._exchange)

  def _get_output_csv_path(self, symbol, root_path):
    td_str = self._trading_date.strftime('%Y%m%d')
    dir_path = os.path.join(root_path,
                            self._machine,
                            td_str,
                            self._exchange,
                            self._owner,
                            self._market_type)
    if not os.path.isdir(dir_path):
      os.makedirs(dir_path, exist_ok=True)
    file_path = '%s/%s.csv' % (dir_path, symbol)
    return file_path


def launch(trading_date, root_dir, machine, market_type, exchange, owner, order_event_dir):
  # 24 will omit some logs..
  hours = 25
  strategy = OrderEventDumper(market_type, exchange, trading_date, machine, owner)
  run_from_og_log_archive(on_log_callback=strategy.on_log2,
                          market_type=market_type,
                          exchange=exchange,
                          start_time=trading_date,
                          end_time=trading_date + datetime.timedelta(hours=hours),
                          root_dir=root_dir,
                          machine=machine,
                          owner=owner)

  strategy.generate_fill_csv(order_event_dir)
  return 0


def main(_):
  start_date = FLAGS.start_date
  assert start_date, '--start_date must be specified.'
  end_date = FLAGS.end_date
  assert end_date, '--end_date must be specified.'

  start_date = convert_string_to_dates(start_date)[0]
  end_date = convert_string_to_dates(end_date)[0]
  trading_dates = [td for td in iterate_date(start_date, end_date)]

  out_root = FLAGS.out_root
  assert out_root, '--out_root must be specified.'

  root_dir = os.path.expanduser('/remote/iosg/strat-1/buckets/log.raw.coin/live/og_proto_log/')

  account_requests = get_all_requests(account_group='Hashed')
  hashed_accounts = set()
  for account_request in account_requests:
    hashed_accounts.add(account_request.owner)

  with ProcessPoolExecutor(max_workers=16) as executor:
    for trading_date in trading_dates:
      td_str = trading_date.strftime('%Y%m%d')
      og_infos = enumerate_og_log_into_oginfos(root_dir, td_str, parse_file_name)
      trading_datetime = datetime.datetime.combine(trading_date, datetime.time.min)
      for og_info in og_infos:
        if og_info.owner not in hashed_accounts:
          continue
        executor.submit(launch,
                        trading_datetime,
                        root_dir,
                        og_info.machine,
                        og_info.market_type,
                        og_info.exchange,
                        og_info.owner,
                        out_root)
        # launch(trading_datetime, root_dir, og_info.machine, og_info.market_type,
        #        og_info.exchange, og_info.owner, out_root)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('start_date', None, 'YYYYmmdd')

  flags.DEFINE_string('end_date', None, 'YYYYmmdd. end_date is not included.')

  flags.DEFINE_string('out_root', None, 'Output root dir.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
